# -*- coding: utf-8 -*-
# @Author  : CaoHan
# @Time    : 2023/11/14 9:02
import json
import math

from ds.uav import UAV, BatteryJ
import ds.uav as uavfile
from ds.sensor import Sensor
import ds.sensor as sensorfile
import ds.container as containerfile
from ds.enumClass import SensorType, ContainerState, Task, ReqType, OffloadDecision
from ds.container import Container, AppMem, AppFlops
from algorithm.offloading import local, ears, score, offloadAll, randomOffload, faasHouse, ecsoa
from generator.harvest import getInterpFunc
from ds.req import COLLAR_INTERVAL, BARN_INTERVAL, BASE_INTERVAL
from generator.arrival import getFixArrival_150_40, getFixArrival_150_30, getFixArrival_150_50
from generator.position import getXY
import heapq

# 面积 1000 * 1000
uavList = []
sensorList = []
jobList = []
containerList = []

funcId = 0
BARN_NUM = 3
BASE_NUM = 3
COW_NUM = 150
reuseTimeWindow = 200

# 结束发射的时间(不能大于18000)
endTransTime = 17000

offloadDecision = OffloadDecision.ECSOA
totalLatency = 0
offloadingNum = [0]
coldStartNum = [0]
avgLatency = [0]
avgOnTimeRatio = [1]
totalReqNum = 0

# plot energy
plot_x = []  # time
remain_energy_list = []  # uav剩余电量之和
consume_energy_list = []  # uav总消耗电量之和
rejection_list = []  # 总rej之和
coldStart_list = []
offloading_list = []
remain_FLOPS_list = []
idle_container_list = []
data_dict = {}
uav_rej_list = [[], [], [], [], [], []]
uav_idleCNum_list = [[], [], [], [], [], []]
uav_remainE_list = [[], [], [], [], [], []]
uav_remainFLOPS_list = [[], [], [], [], [], []]
INTERVAL = 10
interval_num = 0
# 定义能耗插值函数
interp_func = None
# 上一次该函数被卸载的UAV Id
lastOffloadUavId = {}


def init():
    global funcId, uavList, sensorList, jobList, containerList, funcId, totalLatency, offloadingNum, plot_x, \
        remain_energy_list, consume_energy_list, data_dict, interval_num, uav_remainE_list, interp_func, \
        uav_remainFLOPS_list, totalReqNum, rejection_list, coldStartNum, uav_rej_list, uav_idleCNum_list, \
        lastOffloadUavId, coldStart_list, offloading_list, avgLatency, avgOnTimeRatio, remain_FLOPS_list, idle_container_list
    # 重新初始化
    uavList = []
    sensorList = []
    jobList = []
    containerList = []
    funcId = 0
    totalLatency = 0
    uavfile.uavId = 0
    sensorfile.sensorId = 0
    containerfile.containerId = 0
    offloadingNum = [0]
    coldStartNum = [0]
    plot_x = []
    remain_energy_list = []
    consume_energy_list = []
    rejection_list = []
    coldStart_list = []
    offloading_list = []
    remain_FLOPS_list = []
    idle_container_list = []
    avgLatency = [0]
    avgOnTimeRatio = [1]
    uav_rej_list = [[], [], [], [], [], []]
    uav_remainE_list = [[], [], [], [], [], []]
    uav_remainFLOPS_list = [[], [], [], [], [], []]
    uav_idleCNum_list = [[], [], [], [], [], []]
    data_dict = {
        'plot_x': plot_x,
        'remain_energy_list': remain_energy_list,
        'consume_energy_list': consume_energy_list,
        'rejection_list': rejection_list,
        'uav_remainE_list': uav_remainE_list,
        'uav_remainFLOPS_list': uav_remainFLOPS_list,
        'uav_rej_list': uav_rej_list,
        'uav_idleCNum_list': uav_idleCNum_list,
        'offloadingNum': offloadingNum,
        'coldStartNum': coldStartNum,
        'coldStart_list': coldStart_list,
        'offloading_list': offloading_list,
        'avgLatency': avgLatency,
        'avgOnTimeRatio': avgOnTimeRatio,
        'remain_FLOPS_list': remain_FLOPS_list,
        'idle_container_list': idle_container_list
    }
    interval_num = 0
    interp_func = getInterpFunc()
    totalReqNum = 0
    lastOffloadUavId = {}

    # 必须先初始化uav，再sensor
    uavList.extend(
        [UAV(250, 167), UAV(250, 500), UAV(250, 1000 - 167), UAV(750, 167), UAV(750, 500), UAV(750, 1000 - 167)])
    for i in range(BARN_NUM):
        sensorList.append(Sensor(i, SensorType.BARN, i * (BARN_INTERVAL / BARN_NUM)))  # 均匀到达
    for i in range(BASE_NUM):
        sensorList.append(Sensor(i + BARN_NUM, SensorType.BASE, i * (BASE_INTERVAL / BASE_NUM)))  # 均匀到达
    # 固定的配置
    x, y = getXY()
    arrival = getFixArrival_150_40()
    for i in range(COW_NUM):
        sensorList.append(Sensor.initCollar(x[i], y[i], uavList, arrival[i]))

    for sensor in sensorList:
        for i in range(sensor.reqNumEachTrans):
            lastOffloadUavId[(sensor.id, sensor.type.value + 1 + i)] = -1

    # 初始化 jobList req
    for i in range(len(sensorList)):
        sensor = sensorList[i]
        while sensor.transTime <= endTransTime:
            reqList, id = sensor.transmitData(sensor.transTime, funcId)
            funcId = id
            if sensor.type == SensorType.COLLAR:
                # uploadUavId = sensor.move(uavList)
                uploadUavId = sensor.fixMove()  # 假设uav不动！
            for j in range(sensor.reqNumEachTrans):
                req = reqList[j]
                if sensor.type == SensorType.COLLAR:
                    jobList.append((req.start_timestamp, req.funcId, uploadUavId, Task.HANDLE_REQ, req))
                else:
                    jobList.append((req.start_timestamp, req.funcId, sensor.uploadUavId, Task.HANDLE_REQ, req))
            totalReqNum += sensor.reqNumEachTrans
            sensor.transTime += sensor.getInterval()
        # track sensor move
        # print(sensor.trackUavIdList, end=",\n")
    heapq.heapify(jobList)


def containerReuse(req, uavId, time, container):
    # 1. 删除job队列中的CONTAINER_KILL
    killIndex = -1
    for ii in range(len(jobList)):
        if jobList[ii][3] == Task.CON_KILL and jobList[ii][4].containerId == container.id:
            killIndex = ii
            break
    jobList[killIndex] = jobList[-1]
    jobList.pop()
    if killIndex < len(jobList):
        heapq._siftup(jobList, killIndex)
        heapq._siftdown(jobList, 0, killIndex)
    # 2.将req装入容器中
    req.containerId = container.id
    container.run(time)
    heapq.heappush(jobList, (time + req.flop / container.flops, req.funcId, uavId, Task.CON_SPARE, req))
    # 3.删除空闲container id
    uav = uavList[uavId]
    uav.spareConIdList.remove(container.id)


def containerCreation(req, uavId, time):
    uav = uavList[uavId]
    if uav.remainMem < AppMem or uav.remainFLOPS < AppFlops:
        req.isRejected = True
        uav.rejectNum += 1
        return
    # 1. create
    container = Container(uavId, time)
    containerList.append(container)
    req.containerId = container.id
    # 2. container -> uav
    uav.containerIdList.append(container.id)
    uav.remainFLOPS -= container.flops
    uav.remainMem -= container.mem
    # 3. joblist
    heapq.heappush(jobList, (time + container.coldStartTime, req.funcId, uavId, Task.CON_RUN, req))
    coldStartNum[0] += 1


def handleReq(req, uavId, time, solution):
    isOffload, toUavId, toContainerId = decideOffloading(req, uavId, time, solution)
    if isOffload:  # offload
        fromUav = uavList[uavId]
        toUav = uavList[toUavId]
        distance = math.sqrt((fromUav.x - toUav.x) ** 2 + (fromUav.y - toUav.y) ** 2)
        t = fromUav.transT_E(req.size, distance)
        arrivalTime = time + t
        if toContainerId >= 0:  # reuse
            containerReuse(req, toUavId, arrivalTime, containerList[toContainerId])
        else:
            containerCreation(req, toUavId, arrivalTime)
        offloadingNum[0] += 1
    else:  # local
        if toUavId == -200:  # reject
            req.isRejected = True
            uavList[uavId].rejectNum += 1
            return
        # if reuse
        spareConIdList = uavList[uavId].spareConIdList
        if len(spareConIdList) > 0:
            containerReuse(req, uavId, time, containerList[spareConIdList[0]])
            return
        # if not reuse, create container
        containerCreation(req, uavId, time)
    lastOffloadUavId[(req.fromSensorId, req.reqType.value)] = toUavId


def decideOffloading(req, uavId, time, solution):
    if offloadDecision == OffloadDecision.LOCAL:
        return local(req, uavId, time, uavList, containerList)
    elif offloadDecision == OffloadDecision.EARS:
        return ears(req, uavId, time, uavList, containerList, lastOffloadUavId)
    elif offloadDecision == OffloadDecision.SCORE:
        return score(req, uavId, time, uavList, containerList)
    elif offloadDecision == offloadDecision.OFFLOAD_ALL:
        return offloadAll(req, uavId, time, uavList, containerList)
    elif offloadDecision == offloadDecision.RANDOM:
        return randomOffload(req, uavId, time, uavList, containerList)
    elif offloadDecision == offloadDecision.FAAS_HOUSE:
        return faasHouse(req, uavId, time, uavList, containerList, lastOffloadUavId)
    elif offloadDecision == offloadDecision.ECSOA:
        return ecsoa(req, uavId, time, uavList, containerList, 2.143, 13.281)
    elif offloadDecision == offloadDecision.ECSOA_NSGA2:
        return ecsoa(req, uavId, time, uavList, containerList, solution[0], solution[1])


def containerRun(req, uavId, time):
    container = containerList[req.containerId]
    container.run(time)
    heapq.heappush(jobList, (time + req.flop / container.flops, req.funcId, uavId, Task.CON_SPARE, req))


def containerSpare(req, uavId, time):
    if req.containerId not in uavList[uavId].containerIdList:  # UAV之前被关闭过
        return
    global totalLatency
    container = containerList[req.containerId]
    req.end_timestamp = time
    totalLatency += req.end_timestamp - req.start_timestamp
    container.spare(time, reuseTimeWindow)
    heapq.heappush(jobList, (time + reuseTimeWindow, req.funcId, uavId, Task.CON_KILL, req))
    uavList[uavId].spareConIdList.append(container.id)


def containerKill(req, uavId, time):
    if req.containerId not in uavList[uavId].containerIdList:  # UAV之前被关闭过
        return
    container = containerList[req.containerId]
    container.kill()
    uav = uavList[uavId]
    uav.containerIdList.remove(container.id)
    uav.spareConIdList.remove(container.id)
    uav.remainFLOPS += container.flops
    uav.remainMem += container.mem


def sim(solution):
    lastTime = 0
    while len(jobList) != 0:
        (time, funcId, uavId, task, req) = heapq.heappop(jobList)
        # 在任务执行前更新能耗，每次决策时都是最新数据
        if time != lastTime:
            updateEnergyAndData(time)
            lastTime = time
        if uavList[uavId].isOn:
            if task == Task.HANDLE_REQ:
                handleReq(req, uavId, time, solution)
            elif task == Task.CON_RUN:
                containerRun(req, uavId, time)
            elif task == Task.CON_SPARE:
                containerSpare(req, uavId, time)
            elif task == Task.CON_KILL:
                containerKill(req, uavId, time)
            else:
                print("出错了，没找到相应的类型")
        else:
            if task == Task.HANDLE_REQ or task == Task.CON_RUN:
                req.isRejected = True
                uavList[uavId].rejectNum += 1
    if offloadDecision != OffloadDecision.ECSOA_NSGA2:
        for uav in uavList:
            print("id = " + str(uav.id) + " ; rejectNum = " + str(uav.rejectNum) + " ; remain energy = " + str(
                uav.remainEnergy))
        print("总能耗：" + str(getTotalEnergy()))
        print("总拒绝次数：" + str(getTotalRej()))
        print("平均延迟" + str(getAvgLatency()))
        print("机器平均开启时间占比" + str(getAvgOnTimeRatio(lastTime)))
        avgLatency[0] = getAvgLatency()
        avgOnTimeRatio[0] = getAvgOnTimeRatio(lastTime)
        saveData()
        return None, None
    else:
        ret1 = getAvgOnTimeRatio(lastTime)
        ret2 = getTotalRej()
        print(str(solution) + ": 平均开机时间占比：" + str(ret1) + " ;总拒绝次数：" + str(ret2))
        return 1 - ret1, ret2


def saveData():
    # 使用json模块保存数据到JSON文件
    if offloadDecision == OffloadDecision.LOCAL:
        with open('./data/base/LOCAL.json', mode='w') as file:
            json.dump(data_dict, file)
    elif offloadDecision == OffloadDecision.SCORE:
        with open('./data/base/SCORE.json', mode='w') as file:
            json.dump(data_dict, file)
    elif offloadDecision == OffloadDecision.EARS:
        with open('./data/base/EARS.json', mode='w') as file:
            json.dump(data_dict, file)
    elif offloadDecision == OffloadDecision.OFFLOAD_ALL:
        with open('./data/base/OFFLOAD_ALL.json', mode='w') as file:
            json.dump(data_dict, file)
    elif offloadDecision == OffloadDecision.RANDOM:
        with open('./data/base/RANDOM.json', mode='w') as file:
            json.dump(data_dict, file)
    elif offloadDecision == OffloadDecision.FAAS_HOUSE:
        with open('./data/base/FAAS_HOUSE.json', mode='w') as file:
            json.dump(data_dict, file)
    elif offloadDecision == OffloadDecision.ECSOA:
        with open('./data/base/ECSOA.json', mode='w') as file:
            json.dump(data_dict, file)


def updateEnergyAndData(time):
    global interval_num
    remain = 0
    consume = 0
    for uav in uavList:
        uav.updateEnergy(time, containerList, interp_func)
        consume += uav.consumeE
        remain += uav.remainEnergy
    # 给横纵坐标加数据
    interval_num += 1
    if interval_num % INTERVAL == 0:
        remain_FLOPS = 0
        idle_container = 0
        for uav in uavList:
            uav_rej_list[uav.id].append(uav.rejectNum)
            uav_remainE_list[uav.id].append(uav.remainEnergy)
            uav_remainFLOPS_list[uav.id].append(uav.remainFLOPS)
            uav_idleCNum_list[uav.id].append(len(uav.spareConIdList))
            remain_FLOPS += uav.remainFLOPS
            idle_container += len(uav.spareConIdList)
        plot_x.append(time)
        remain_energy_list.append(remain)
        consume_energy_list.append(consume)
        rejection_list.append(getTotalRej())
        coldStart_list.append(coldStartNum[0])
        offloading_list.append(offloadingNum[0])
        remain_FLOPS_list.append(remain_FLOPS)
        idle_container_list.append(idle_container)


def getTotalEnergy():
    return consume_energy_list[-1]


def getTotalRej():
    r = 0
    for uav in uavList:
        r += uav.rejectNum
    return r


def getAvgLatency():
    return totalLatency / (totalReqNum - getTotalRej())


def getAvgOnTimeRatio(lastTime):  # 机器平均开启时间占比
    close_sum = 0
    for uav in uavList:
        close_sum += uav.closeTime
    return 1 - close_sum / (len(uavList) * lastTime)


def simulation(solution):
    init()
    if offloadDecision == OffloadDecision.ECSOA_NSGA2:
        return sim(solution)
    else:
        sim(None)


if __name__ == "__main__":
    print("start...")
    simulation(None)
